#!/usr/bin/env python3

import rospy
import time
from std_msgs.msg import String
from std_msgs.msg import Int32
from mavproxy_ros1.msg import FlightCmd
from mavproxy_ros1.msg import MavproxyState
from geometry_msgs.msg import PoseStamped
from mavproxy_tool import MavproxyTool

mavproxy_state = MavproxyState()
drone_pose = PoseStamped()

STABILIZE_PX4_MODE = (0, "STABILIZED")
TAKEOFF_PX4_MODE = (1, "AUTO.TAKEOFF")
LAND_PX4_MODE = (2, "AUTO.LAND")
HOLD_PX4_MODE = (3, "AUTO.LOITER")
RETURN_PX4_MODE = (4, "RTL")
OFFBOARD_PX4_MODE = (5, "OFFBOARD")
MISSION_PX4_MODE = (6, "MISSION")

def state_callback(state_msg):
    global mavproxy_state
    mavproxy_state = state_msg

def position_callback(position_msg):
    global drone_pose
    drone_pose = position_msg

def handle_takeoff(command_full):
    command_type = command_full.split(" ")[0]
    rospy.loginfo(command_type)

    if mavproxy_state.flight_mode == STABILIZE_PX4_MODE[1] or mavproxy_state.flight_mode == HOLD_PX4_MODE[1]:
        command_pub.publish(tool.get_flight_cmd("flight_arm"))
        time.sleep(2)
        command_pub.publish(tool.get_flight_cmd(f"flight_mode {TAKEOFF_PX4_MODE[0]}"))
    else:
        rospy.loginfo(mavproxy_state.flight_mode)
        
        command_pub.publish(tool.get_flight_cmd(f"flight_mode {STABILIZE_PX4_MODE[0]}"))
        time.sleep(1)

        if (mavproxy_state.flight_mode == STABILIZE_PX4_MODE[1]):
            command_pub.publish(tool.get_flight_cmd("flight_arm"))
            time.sleep(2)
            command_pub.publish(tool.get_flight_cmd(f"flight_mode {TAKEOFF_PX4_MODE[0]}"))
        else:
            rospy.logerr("[parser] Failed to takeoff because failed to change the mode")
            return

def handle_land(command_full):
    command_type = command_full.split(" ")[0]
    rospy.loginfo(command_type)

    if mavproxy_state.flight_mode == HOLD_PX4_MODE[1]:
        command_pub.publish(tool.get_flight_cmd(f"flight_mode {LAND_PX4_MODE[0]}"))
    else:
        command_pub.publish(tool.get_flight_cmd(f"flight_mode {HOLD_PX4_MODE[0]}"))
        time.sleep(2)
        command_pub.publish(tool.get_flight_cmd(f"flight_mode {LAND_PX4_MODE[0]}"))

def handle_return(command_full):
    command_type = command_full.split(" ")[0]
    rospy.loginfo(command_type)

    if mavproxy_state.flight_mode == HOLD_PX4_MODE[1]:
        command_pub.publish(tool.get_flight_cmd("flight_return"))

def handle_hover(command_full):
    command_type = command_full.split(" ")[0]
    rospy.loginfo(command_type)

    if mavproxy_state.flight_mode != HOLD_PX4_MODE[1]:
        command_pub.publish(tool.get_flight_cmd("flight_hover"))

def handle_fly_target_gpos(command_full):
    command_type = command_full.split(" ")[0]
    rospy.loginfo(command_type)

    command_pub.publish(tool.get_flight_cmd(f"flight_stop"))
    time.sleep(1)

    command_pub.publish(tool.get_flight_cmd(command_full))

# 飞往指定位置：向前、后、左、右、上、下
def handle_fly_target_pos(command_full):
    command_type = command_full.split(" ")[0]
    rospy.loginfo(command_type)

    command_pub.publish(tool.get_flight_cmd(f"flight_stop"))
    time.sleep(1)

    target_position = [
        drone_pose.pose.position.x,
        drone_pose.pose.position.y,
        drone_pose.pose.position.z,
        ]

    if mavproxy_state.flight_mode != HOLD_PX4_MODE[1] and mavproxy_state.flight_mode != OFFBOARD_PX4_MODE[1]:
        command_pub.publish(tool.get_flight_cmd(f"flight_mode {HOLD_PX4_MODE[0]}"))

        time.sleep(1)
        if mavproxy_state.flight_mode != HOLD_PX4_MODE[1]:
            rospy.logerr("[parser] Failed to change mode to hold")
            return

    if command_type == "flight_targetpos":
        target_position[0] = float(command_full.split(" ")[1])
        target_position[1] = float(command_full.split(" ")[2])
        target_position[2] = float(command_full.split(" ")[3])
    elif command_type == "flight_forward":
        param = float(command_full.split(" ")[1])
        target_position[0] += param
        pass
    elif command_type == "flight_back":
        param = float(command_full.split(" ")[1])
        target_position[0] -= param
        pass
    elif command_type == "flight_left":
        param = float(command_full.split(" ")[1])
        target_position[1] += param
        pass
    elif command_type == "flight_right":
        param = float(command_full.split(" ")[1])
        target_position[1] -= param
        pass
    elif command_type == "flight_up":
        param = float(command_full.split(" ")[1])
        target_position[2] += param
        pass
    elif command_type == "flight_down":
        param = float(command_full.split(" ")[1])
        target_position[2] -= param 
        pass

    rospy.loginfo(f"[TD DEBUG][parser] local: x {drone_pose.pose.position.x}, y {drone_pose.pose.position.y}, z {drone_pose.pose.position.z}")
    rospy.loginfo(f"[TD DEBUG][parser] target: x {target_position[0]}, y {target_position[1]}, z {target_position[2]}")

    command_pub.publish(tool.get_flight_cmd(f"flight_targetpos {target_position[0]} {target_position[1]} {target_position[2]}"))

def handle_default(command_full):
    command_type = command_full.split(" ")[0]
    rospy.loginfo(f"[parser] Cmd {command_type} use default handle")
    command_pub.publish(tool.get_flight_cmd(command_full))

def handle_help(command_full):
    for cmd_name in HANDLE_FUNC_MAP.keys():
        rospy.loginfo(cmd_name)

HANDLE_FUNC_MAP = {
    "flight_arm": handle_default,
    "flight_disarm": handle_default,
    "flight_takeoff": handle_takeoff,
    "flight_land": handle_land,
    "flight_return": handle_return,
    "flight_hover": handle_hover,
    "flight_forward": handle_fly_target_pos,
    "flight_back": handle_fly_target_pos,
    "flight_left": handle_fly_target_pos,
    "flight_right": handle_fly_target_pos,
    "flight_up": handle_fly_target_pos,
    "flight_down": handle_fly_target_pos,
    "flight_targetpos": handle_fly_target_pos,
    "flight_mission": handle_fly_target_pos,
    "flight_targetgpos": handle_fly_target_gpos,
    "flight_mode": handle_default,
    "flight_stop": handle_default,
    "help": handle_help,
}

def get_handle_func(input_msg):
    handle_func = HANDLE_FUNC_MAP.get(input_msg.data.split(" ")[0])
    if handle_func:  
        handle_func(input_msg.data)
    else:  
        rospy.logwarn(f"[parser] Unknown command type, cmd is {input_msg.data}")

if __name__ == '__main__':  
    try:
        rospy.init_node('flight_cmd_parser', anonymous=True)
        tool = MavproxyTool()
        input_sub = rospy.Subscriber('/mavproxy_ros1/input_cmd', String, get_handle_func, queue_size=10)
        command_pub = rospy.Publisher('/mavproxy_ros1/parsed_cmd', FlightCmd, queue_size=10)
        state_sub = rospy.Subscriber('/mavproxy_ros1/state', MavproxyState, state_callback, queue_size=10)
        local_position_sub = rospy.Subscriber('mavros/local_position/pose', PoseStamped, position_callback, queue_size=10)
        rospy.spin()
    except rospy.ROSInterruptException:
        pass
